import os
from typing import Dict, List
import httpx
import dashscope
from mcp.server.fastmcp import FastMCP


def get_zhipuai_api_key() -> str:
    """获取环境变量中的智谱AI API密钥"""
    api_key = os.getenv("ZHIPUAI_API_KEY")
    if not api_key:
        raise ValueError("ZHIPUAI_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


def get_aliyun_api_key() -> str:
    """获取环境变量中的阿里云API密钥"""
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("DASHSCOPE_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


def format_search_results(search_data: Dict) -> str:
    """
    格式化搜索结果数据

    参数:
        search_data: 搜索结果原始数据

    返回:
        格式化后的搜索结果字符串
    """
    formatted = ""

    # 添加搜索意图信息
    if "search_intent" in search_data and search_data["search_intent"]:
        formatted += "搜索意图分析:\n"
        for intent in search_data["search_intent"]:
            formatted += f"- 查询: {intent.get('query', 'N/A')}\n"
            formatted += f"  意图: {intent.get('intent', 'N/A')}\n"
            formatted += f"  关键词: {intent.get('keywords', 'N/A')}\n"
        formatted += "\n"

    # 添加搜索结果
    if "search_result" in search_data and search_data["search_result"]:
        formatted += "搜索结果:\n"
        formatted += "-" * 50 + "\n"

        for i, result in enumerate(search_data["search_result"], 1):
            formatted += f"{i}. 标题: {result.get('title', 'N/A')}\n"
            formatted += f"   链接: {result.get('link', 'N/A')}\n"
            formatted += f"   内容: {result.get('content', 'N/A')}\n"

            # 添加可选字段
            if result.get("media"):
                formatted += f"   媒体: {result['media']}\n"
            if result.get("publish_date"):
                formatted += f"   发布日期: {result['publish_date']}\n"

            formatted += "-" * 50 + "\n"
    else:
        formatted = "未找到相关搜索结果"

    return formatted


def format_qwen_response(response_data) -> str:
    """
    格式化Qwen模型的响应数据

    参数:
        response_data: Qwen模型的响应数据

    返回:
        格式化后的响应字符串
    """
    formatted = ""
    
    # 添加搜索信息（如果存在）
    if hasattr(response_data.output, 'search_info') and response_data.output.search_info:
        search_results = response_data.output.search_info.get("search_results", [])
        if search_results:
            formatted += "=" * 20 + "搜索信息" + "=" * 20 + "\n"
            for web in search_results:
                formatted += f"[{web['index']}]: [{web['title']}]({web['url']})\n"
            formatted += "\n"

    # 添加回复内容
    formatted += "=" * 20 + "回复内容" + "=" * 20 + "\n"
    if hasattr(response_data.output, 'choices') and response_data.output.choices:
        content = response_data.output.choices[0].message.content
        formatted += content
    else:
        formatted += "无内容"

    # 添加使用统计信息
    if hasattr(response_data, 'usage'):
        usage = response_data.usage
        formatted += f"\n\n---\nToken使用情况: 输入{usage.input_tokens}，输出{usage.output_tokens}，总计{usage.total_tokens}"

    return formatted


mcp = FastMCP("Web Search MCP Server")


@mcp.tool()
async def web_search_zhipu(
    query: str, search_engine: str = "search_std", search_intent: bool = False
) -> str:
    """
    执行网络搜索

    参数:
        query: 搜索查询词
        search_engine: 搜索引擎类型，可选值:
                      "search_std" - 智谱基础版搜索引擎
                      "search_pro" - 智谱高阶版搜索引擎
                      "search_pro_sogou" - 搜狗
                      "search_pro_quark" - 夸克搜索
                      默认值为"search_std"
        search_intent: 是否返回搜索意图分析，默认为False

    返回:
        格式化后的搜索结果字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        httpx.RequestError: 网络请求失败
    """
    # 验证参数
    valid_engines = ["search_std", "search_pro", "search_pro_sogou", "search_pro_quark"]
    if search_engine not in valid_engines:
        raise ValueError(f"无效的搜索引擎类型。支持的类型: {valid_engines}")

    BIGMODEL_SEARCH_URL = "https://open.bigmodel.cn/api/paas/v4/web_search"

    api_key = get_zhipuai_api_key()
    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {
        "search_query": query,
        "search_engine": search_engine,
        "search_intent": search_intent,
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(
                BIGMODEL_SEARCH_URL, headers=headers, json=payload
            )
            response.raise_for_status()
            data = response.json()

            return format_search_results(data)
        except httpx.RequestError as e:
            raise httpx.RequestError(f"网络请求失败: {str(e)}")
        except Exception as e:
            raise ValueError(f"搜索处理错误: {str(e)}")


@mcp.tool()
async def web_search_qwen(
    query: str,
    model: str = "qwen-plus",
    forced_search: bool = True,
    search_strategy: str = "turbo",
    enable_source: bool = True,
    enable_citation: bool = True,
    citation_format: str = "[ref_<number>]",
) -> str:
    """
    使用阿里云Qwen模型执行网络搜索

    参数:
        query: 搜索查询词
        model: Qwen模型版本，可选值: "qwen-turbo"、"qwen-plus"、"qwen-max"等
        forced_search: 是否强制联网搜索，True表示强制开启，False表示不强制开启
        search_strategy: 搜索策略，可选值:
                        "turbo" - 兼顾响应速度与搜索效果（默认值）
                        "max" - 高性能模式，提供最优效果
        enable_source: 是否返回搜索来源信息
        enable_citation: 是否开启角标标注功能
        citation_format: 角标样式，可选值:
                        "[<number>]" - 角标形式为[i]
                        "[ref_<number>]" - 角标形式为[ref_i]

    返回:
        格式化后的模型响应字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        httpx.RequestError: 网络请求失败
    """
    # 验证参数
    valid_models = ["qwen-turbo", "qwen-plus", "qwen-max"]
    if model not in valid_models:
        raise ValueError(f"无效的模型类型。支持的类型: {valid_models}")

    valid_strategies = ["turbo", "max"]
    if search_strategy not in valid_strategies:
        raise ValueError(f"无效的搜索策略。支持的策略: {valid_strategies}")
        
    valid_citation_formats = ["[<number>]", "[ref_<number>]"]
    if citation_format not in valid_citation_formats:
        raise ValueError(f"无效的角标样式。支持的样式: {valid_citation_formats}")

    api_key = get_aliyun_api_key()
    dashscope.api_key = api_key

    try:
        response = dashscope.Generation.call(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": query},
            ],
            enable_search=True,
            search_options={
                "forced_search": forced_search,
                "enable_source": enable_source,
                "enable_citation": enable_citation,
                "citation_format": citation_format,
                "search_strategy": search_strategy,
            },
            result_format="message",
        )

        return format_qwen_response(response)
    except Exception as e:
        raise ValueError(f"搜索处理错误: {str(e)}")


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()